﻿using IconExtractor;
using System;
using System.Collections.Concurrent;
using System.Drawing;
using System.Drawing.Imaging;
using System.IO;
using System.Runtime.InteropServices;
using System.Windows;
using System.Windows.Interop;
using System.Windows.Media;
using System.Windows.Media.Imaging;

namespace ScriptBox
{
    // 在文件顶部添加新的using
    using System.Drawing;
    using System.Drawing.Imaging;
    
    // 在类声明前添加特性
    [System.Runtime.Versioning.SupportedOSPlatform("windows")]
    class ExeIcon : IDisposable
    {
        private static class NativeMethods
        {
            [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Auto)]
            public struct SHFILEINFO
            {
                public IntPtr hIcon;
                public int iIcon;
                public uint dwAttributes;
                [MarshalAs(UnmanagedType.ByValTStr, SizeConst = 260)]
                public string szDisplayName;
                [MarshalAs(UnmanagedType.ByValTStr, SizeConst = 80)]
                public string szTypeName;
            }

            [DllImport("shell32.dll", CharSet = CharSet.Auto)]
            public static extern IntPtr SHGetFileInfo(string pszPath, uint dwFileAttributes,
                ref SHFILEINFO psfi, uint cbSizeFileInfo, uint uFlags);

            [DllImport("user32.dll")]
            public static extern bool DestroyIcon(IntPtr hIcon);

            [DllImport("shell32.dll", CharSet = CharSet.Auto)]
            public static extern int ExtractIconEx(string lpszFile, int nIconIndex, 
                out IntPtr phiconLarge, out IntPtr phiconSmall, int nIcons);
        }

        // 修改缓存结构并添加过期时间
        private static readonly ConcurrentDictionary<string, (ImageSource image, DateTime timestamp)> _imageCache = new();
        private static readonly TimeSpan CacheExpiration = TimeSpan.FromMinutes(30);
        private bool _disposed;

        public static ImageSource GetIconImage(string exePath)
        {
            if (!File.Exists(exePath)) return null;

            // 检查图片缓存
            if (_imageCache.TryGetValue(exePath, out var cacheEntry) && 
                DateTime.Now - cacheEntry.timestamp < CacheExpiration)
            {
                return cacheEntry.image;
            }

            try
            {
                using var icon = GetBestIcon(exePath);
                var image = CreateImageSource(icon);
                
                // 缓存冻结后的ImageSource
                image.Freeze();
                _imageCache.AddOrUpdate(exePath, 
                    (image, DateTime.Now), 
                    (_, _) => (image, DateTime.Now));
                
                return image;
            }
            catch
            {
                return CreateImageSource(GetFallbackIcon());
            }
        }

        private static Icon GetBestIcon(string exePath)
        {
            var resultIcon = GetIcon_With_IconExtractor(exePath, 96) 
                ?? GetIcon_With_ExtractIconEx(exePath)
                ?? GetIcon(exePath)
                ?? GetFallbackIcon();
            return resultIcon;
        }

        private static ImageSource CreateImageSource(Icon icon)
        {
            if (icon == null) return null;

            var image = Imaging.CreateBitmapSourceFromHIcon(
                icon.Handle,
                Int32Rect.Empty,
                BitmapSizeOptions.FromWidthAndHeight(icon.Width, icon.Height));
            
            image.Freeze();
            return image;
        }

        public void Dispose()
        {
            if (_disposed) return;
            
            // 清理图标缓存
            _imageCache.Clear();
            
            _disposed = true;
            GC.SuppressFinalize(this);
        }

        private static Icon GetFallbackIcon()
        {
            return SystemIcons.Application;
        }

        /// <summary>
        /// 通过ExtractIconEx方法获取，实测大图标尺寸32，小图标尺寸16
        /// </summary>
        /// <param name="exePath"></param>
        /// <returns></returns>
        public static Icon GetIcon_With_ExtractIconEx(string exePath)
        {
            IntPtr largeIconHandle, smallIconHandle;
            int iconCount = NativeMethods.ExtractIconEx(exePath, 0, 
                out largeIconHandle, out smallIconHandle, 1);
            
            if (largeIconHandle != IntPtr.Zero)
            {
                var icon = (Icon)Icon.FromHandle(largeIconHandle).Clone();
                NativeMethods.DestroyIcon(largeIconHandle);
                return icon;
            }
            return null;
        }

        /// <summary>
        /// 通过Icon类自有方法获取，实测图标尺寸32
        /// </summary>
        /// <param name="exePath"></param>
        /// <returns></returns>
        public static Icon GetIcon(string exePath)
        {
            if (!File.Exists(exePath)) return null;

            try
            {
                using var _ = new Bitmap(1, 1);
                return Icon.ExtractAssociatedIcon(exePath);
            }
            catch
            {
                return GetFallbackIcon();
            }
        }
        /// <summary>
        /// 这种方法可以提取大尺寸图标，但是如果exe中不存在大尺寸，会被强行拉伸。
        /// </summary>
        /// <param name="exePath"></param>
        /// <param name="size">需要获取的图标尺寸</param>
        /// <returns>返回能找到的最大尺寸图标</returns>
        public static Icon GetIcon_With_IconExtractor(string exePath, int size)
        {
            uint _nIcons = PInvoke.PrivateExtractIcons(exePath, 0, 0, 0, null, null, 0, 0);
            if (_nIcons == 0) return null;

            IntPtr[] phicon = new IntPtr[_nIcons];
            uint[] piconid = new uint[_nIcons];
            uint nIcons = PInvoke.PrivateExtractIcons(exePath, 0, size, size, phicon, piconid, _nIcons, 0);
            
            if (nIcons > 0 && phicon[0] != IntPtr.Zero)
            {
                var icon = (Icon)Icon.FromHandle(phicon[0]).Clone();
                NativeMethods.DestroyIcon(phicon[0]);
                return icon;
            }
            return null;
        }
    }
}
